package com.lw.leetcode.sort;

import java.util.*;

/**
 * 692. 前K个高频单词
 *
 * @Author liw
 * @Date 2021/5/20 12:38
 * @Version 1.0
 */
public class TopKFrequent {
    public List<String> topKFrequent(String[] words, int k) {
        if (words == null || words.length == 0) {
            return Collections.emptyList();
        }
        int length = words.length;
        Map<String, Integer> map = new HashMap<>(length << 1);
        for (String word : words) {
            map.merge(word, 1, (a, b) -> a + b);
        }
        List<String> list = new ArrayList<>(k);


        TreeMap<Integer, TreeSet<String>> treeMap = new TreeMap<>((a, b) -> b - a);

        for (Map.Entry<String, Integer> entry : map.entrySet()) {
            treeMap.merge(entry.getValue(), new TreeSet<String>() {{
                add(entry.getKey());
            }}, (a, b) -> {
                a.addAll(b);
                return a;
            });
        }

        int n = 0;
        for (Map.Entry<Integer, TreeSet<String>> entry : treeMap.entrySet()) {
            TreeSet<String> value = entry.getValue();
            for (String s : value) {
                list.add(s);
                if (++n == k) {
                    return list;
                }
            }
        }
        return list;
    }


    public List<String> topKFrequent2(String[] words, int k) {
        Map<String, Integer> cnt = new HashMap<>();
        for (String word : words) {
            cnt.put(word, cnt.getOrDefault(word, 0) + 1);
        }
        List<String> rec = new ArrayList<>();
        for (Map.Entry<String, Integer> entry : cnt.entrySet()) {
            rec.add(entry.getKey());
        }
        rec.sort((a, b) -> cnt.get(a).equals(cnt.get(b)) ? a.compareTo(b) : cnt.get(b) - cnt.get(a));
        return rec.subList(0, k);
    }



}
